from core.function import Function
from core.cuda import cuda
import utils


# 线性变换
# 将 y = x.dot(W)+b 封装起来
# 线性变换对应神经网络的全连接层
class Linear(Function):
    def forward(self, x, W, b):
        """
        执行线性变换的前向传播。

        Args:
            x (Variable): 输入变量。
            W (ndarray): 权重参数。
            b (ndarray or None): 偏置参数。

        Returns:
            Variable: 前向传播的输出变量。
        """
        W = cuda.to_array(W)
        x = cuda.to_array(x)
        y = x.dot(W)
        if b is not None:
            b = cuda.to_array(b)
            y = y + b
        return y

    def backward(self, gy):
        """
        执行线性变换的反向传播。

        Args:
            gy (Variable): 相对于输出的梯度。

        Returns:
            tuple: 包含相对于输入 x, 权重 W, 偏置 b 的梯度的元组。
        """

        x, W, b = self.inputs
        gb = None if b.data is None else utils.functions_collect.sum_to(gy, b.shape)
        gx = utils.functions_collect.matmul(gy, W.T)
        gW = utils.functions_collect.matmul(x.T, gy)
        return gx, gW, gb

    def linear(self, x, W, b=None):
        """
        执行线性变换的计算。

        Args:
            x (Variable): 输入变量。
            W (ndarray): 权重参数。
            b (ndarray or None): 偏置参数。

        Returns:
            Variable: 输出变量。
        """
        return self(x, W, b)
